{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamics of the Double Pendulum\n",
    "\n",
    "This notebook provides a simple example of simulating the dynamics of a double pendulum, and examining the equations of motion in manipulator form."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YoTJGozDkYU7"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running a simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KN--nx6giWw2"
   },
   "outputs": [],
   "source": [
    "# This cell sets up a double pendulum, runs a simulation, and\n",
    "# renders a video of its results.\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph,\n",
    "                         DiagramBuilder,\n",
    "                         Parser,\n",
    "                         PlanarSceneGraphVisualizer,\n",
    "                         Simulator,\n",
    "                         SignalLogger)\n",
    "from underactuated import FindResource\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "# Set up a block diagram with the robot (dynamics) and a visualization block.\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "\n",
    "# Load the double pendulum from Universal Robot Description Format\n",
    "parser = Parser(plant, scene_graph)\n",
    "parser.AddModelFromFile(FindResource(\"models/double_pendulum.urdf\"))\n",
    "plant.Finalize()\n",
    "\n",
    "builder.ExportInput(plant.get_actuation_input_port())\n",
    "visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph,\n",
    "                                                          xlim=[-2.8, 2.8],\n",
    "                                                          ylim=[-2.8, 2.8],\n",
    "                                                          show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "\n",
    "# The logger needs to be told to expect a 4-element input\n",
    "# (the 4-element double pendulum state, in this case).\n",
    "logger = builder.AddSystem(SignalLogger(4))\n",
    "logger.DeclarePeriodicPublish(0.033333, 0.0)\n",
    "builder.Connect(plant.get_state_output_port(), logger.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# Set up a simulator to run this diagram\n",
    "simulator = Simulator(diagram)\n",
    "\n",
    "# Set the initial conditions\n",
    "context = simulator.get_mutable_context()\n",
    "context.SetContinuousState([1., 1., 0., 0.])  # (theta1, theta2, theta1dot, theta2dot)\n",
    "context.FixInputPort(0, [0.,0.])   # Zero input torques\n",
    "\n",
    "# Simulate and animate\n",
    "AdvanceToAndVisualize(simulator, visualizer, 10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vYkJeGwviWxd"
   },
   "outputs": [],
   "source": [
    "# Run this cell after running a simulation to plot traces\n",
    "# of the pendulum state across time.\n",
    "fields = [\"shoulder\", \"elbow\"]\n",
    "for i in range(2):\n",
    "    #plt.subplot(3, 1, i+1)\n",
    "    plt.plot(logger.sample_times(), logger.data()[(i, i+2), :].transpose())\n",
    "    plt.legend([\"position\", \"velocity\"])\n",
    "    plt.xlabel('t')\n",
    "    plt.ylabel(fields[i])\n",
    "    plt.grid(True)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's worth taking a peek at the [file that describes the robot](https://github.com/RussTedrake/underactuated/blob/master/underactuated/models/double_pendulum.urdf). URDF and SDF are two of the standard formats, and they can be used to describe even very complicated robots (like the Boston Dynamics humanoid).\n",
    "\n",
    "## Inspecting the dynamics (the manipulator equations)\n",
    "\n",
    "We can also use Drake to evaluate the manipulator equations.  First we will evaluate the manipulator equations for a particular robot (with numerical values assigned for mass, link lengths, etc) and for a particular state of the robot.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.all import MultibodyPlant\n",
    "from underactuated import ManipulatorDynamics\n",
    "\n",
    "plant = MultibodyPlant(time_step=0)\n",
    "parser = Parser(plant)\n",
    "parser.AddModelFromFile(FindResource(\"models/double_pendulum.urdf\"))\n",
    "plant.Finalize()\n",
    "\n",
    "# Evaluate the dynamics numerically\n",
    "q = [0.1, 0.1]\n",
    "v = [1, 1]\n",
    "(M, Cv, tauG, B, tauExt) = ManipulatorDynamics(plant, q, v)\n",
    "print(\"M = \\n\" + str(M))\n",
    "print(\"Cv = \" + str(Cv))\n",
    "print(\"tau_G = \" + str(tauG))\n",
    "print(\"B = \" + str(B))\n",
    "print(\"tau_ext = \" + str(tauExt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drake is also fairly unique in supporting symbolic computation (with floating point coefficients).  Here is an example of printing out the symbolic dynamics of the double pendulum.  (If you've ever written out the equations of your robot, you know they get complicated quickly!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.all import Variable\n",
    "\n",
    "# Evaluate the dynamics symbolically\n",
    "q = [Variable(\"theta0\"), Variable(\"theta1\")]\n",
    "v = [Variable(\"thetadot0\"), Variable(\"thetadot1\")]\n",
    "(M, Cv, tauG, B, tauExt) = ManipulatorDynamics(plant.ToSymbolic(), q, v)\n",
    "print(\"M = \\n\" + str(M))\n",
    "print(\"Cv = \" + str(Cv))\n",
    "print(\"tau_G = \" + str(tauG))\n",
    "print(\"B = \" + str(B))\n",
    "print(\"tau_ext = \" + str(tauExt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use a similar mechanisms to support automatic differentiation; we'll see examples of that soon!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "dynamics.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}